#!/usr/bin/env python

import os, sys
import numpy as np
from astropy.io import fits
from astropy.visualization import ZScaleInterval
import matplotlib.pylab as plt
import time

from testlib.image import polyfit


def apply_gain(fimage, fgain):
    img = fits.getdata(fimage)
    gain = fits.getdata(fgain)
    
#    with fits.open(fimage) as img, fits.open(fgain) as gain:
#        data = img[0].data  * gain[0].data
    
    data = img * gain
    return data

#def polyfit_flat(fimage, fout, degree=3, ffig=None):
#    with fits.open(fimage) as img:
#        data = polyfit(img[0].data, degree=degree)
#        hdu = fits.PrimaryHDU(data, header=img[0].header)
#        hdu.header['POLYDEG'] = degree
#        ratio = img[0].data / data
#        hdu2 = fits.ImageHDU(ratio, header=img[0].header)
#        hdu2.header['POLYDEG'] = degree
#        hdulist = fits.HDUList([hdu, hdu2])
#        hdulist.writeto(fout, overwrite=True)

#        if ffig is not None:
#            zs = ZScaleInterval()
#            vmin, vmax = zs.get_limits(img[0].data)
#            plt.figure(figsize=(18, 5))

#            plt.subplot(131)
#            plt.imshow(img[0].data, vmin=vmin, vmax=vmax, cmap='Blues_r', origin='lower')
#            plt.title('wavelength = {:.1f} nm'.format(img[0].header['WAVELEN']))
#            plt.colorbar()

#            plt.subplot(132)
#            plt.imshow(data, vmin=vmin, vmax=vmax, cmap='Blues_r', origin='lower')
#            plt.title('polynomial degree = {}'.format(degree))
#            plt.colorbar()

#            plt.subplot(133)
#            lims = zs.get_limits(ratio)
#            plt.imshow(ratio, vmin=lims[0], vmax=lims[1], cmap='rainbow', origin='lower')
#            plt.colorbar()

#            plt.tight_layout()
#            plt.savefig(ffig)
#            plt.close()

def polyfit_flat(flate, ffig1, ffig2, degree=4):
    """
    return ratio
    """
    data = polyfit(flate, degree=degree)
    ratio = flate / data
    return ratio
    

def plot_prnu(fimage_list, fgain):
    plt.figure(figsize=(12,10))
    for i in range(len(fimage_list)):
        fimage = fimage_list[i]
        
        print('> processing: {}'.format(fimage))
        
        # convert DN to e-
        flate = apply_gain(fimage, fgain )
        
        ffig1 = fimage.replace('.fits', '.png')
        ffig2 = fimage.replace('.fits', '_hist.png')
        wave = fimage.split('_')[1]
        
        # fit large scale structure of the flat
        ratio = polyfit_flat(flate=flate, ffig1=ffig1, ffig2=ffig2)

        rm = np.median(ratio)
        rms = np.std(ratio)
        idx1 = ratio < rm - 5*rms
        idx2 = ratio > rm + 5*rms
        ratio[idx1] = rm
        ratio[idx2] = rm

        ax = plt.subplot(2,2,i+1)
        zs = ZScaleInterval()
        lims = zs.get_limits(ratio)
        lims = [0.915, 1.05]
#        lims = [0.975, 1.05]
        lims = [0.9, 1.05]

        tmp = ratio.flatten()
        idx = np.logical_and(tmp>0.95, tmp<1.05)

#        rms = 'rms: {:.3f}'.format(np.std(tmp))

        plt.imshow(ratio, vmin=lims[0], vmax=lims[1], cmap='Blues_r', origin='lower')
        plt.colorbar()
#        plt.title(r'$\lambda = ' + wave+ ' PRNU: {:.3f}'.format(np.std(tmp)))
        plt.title(r'$\lambda$ = ' + wave)
        

    plt.tight_layout()
    plt.savefig('4x4.png')
    

def get_flist(flist):
    fimage_list = []
    fp = open(flist,'r')
    lines = fp.readlines()
    for line in lines:
        fimage_list.append(line.replace('\n',''))
        
    return fimage_list


def main():
    fimage_list = get_flist(sys.argv[1])
    plot_prnu(fimage_list, sys.argv[2])
    

if __name__ == '__main__':
    main()


    
